import os
import json
import numpy as np
import matplotlib.pyplot as plt

# 解析 JSON 并提取数据
def parse_json(json_path, bandType="R"):
    with open(json_path, "r", encoding="utf-8") as f:
        data = json.load(f)

    time_hours = []
    bg_median_values = []
    bg_std_values = []
    ee70_values = []
    ee70_rms_values = []
    sun_ang_values = []
    moon_ang_values = []

    for img_type in data:
        for img in data[img_type]["images"]:
            date_obs = img.get("DATE-OBS", "")
            if not date_obs:
                continue

            try:
                # 转换时间为小时
                h, m, s = map(int, date_obs.split("T")[1].split(":"))
                time_decimal = h + m / 60 + s / 3600
                time_hours.append(time_decimal)
                
                if bandType == "B":
                    image_info = img.get("image_info_b", {})
                else:
                    image_info = img.get("image_info_r", {})
                
                if len(image_info)>0:
                    bg_median_values.append(image_info.get("BGMEDIAN", np.nan))
                    bg_std_values.append(image_info.get("BGSTD", np.nan))
                    ee70_values.append(image_info.get("EE70", np.nan))
                    ee70_rms_values.append(image_info.get("EE70_RMS", np.nan))
                    sun_ang_values.append(image_info.get("SUN_ANG", np.nan))
                    moon_ang_values.append(image_info.get("MOON_ANG", np.nan))
            except Exception as e:
                print(f"Error processing DATE-OBS {date_obs}: {e}")

    sorted_data = sorted(zip(time_hours, bg_median_values, bg_std_values, ee70_values, ee70_rms_values, sun_ang_values, moon_ang_values))
    if sorted_data:
        time_hours, bg_median_values, bg_std_values, ee70_values, ee70_rms_values, sun_ang_values, moon_ang_values = zip(*sorted_data)
    else:
        return [], [], [], [], [], [], []

    return (
        np.array(time_hours),
        np.array(bg_median_values, dtype=np.float64),
        np.array(bg_std_values, dtype=np.float64),
        np.array(ee70_values, dtype=np.float64),
        np.array(ee70_rms_values, dtype=np.float64),
        np.array(sun_ang_values, dtype=np.float64),
        np.array(moon_ang_values, dtype=np.float64),
    )

# 3σ 过滤函数
def filter_3sigma(data, maxValue):
    valid_data = data[~np.isnan(data)]
    if len(valid_data) < 2:
        return np.full_like(data, np.nan)

    mean = np.nanmean(valid_data)
    std = np.nanstd(valid_data)
    # trst = np.where((data >= mean - sigmaThred * std) & (data <= mean + sigmaThred * std), data, np.nan)
    trst = np.where(data <= maxValue, data, np.nan)
    return trst

# 绘制子图
def plot_data(save_dir, date_str, time_hours, bg_medianr, ee70r, bg_medianb, ee70b, sun_angs, moon_angs):

    os.makedirs(save_dir, exist_ok=True)

    plt.figure(figsize=(12, 9))

    # 计算 3σ 过滤数据
    bg_median_filteredr = filter_3sigma(bg_medianr, 800)
    ee70_filteredr = filter_3sigma(ee70r, 3)
    bg_median_filteredb = filter_3sigma(bg_medianb, 800)
    ee70_filteredb = filter_3sigma(ee70b, 4)
    
    # 定义 X 轴刻度（0 ~ 23，间隔 1 小时）
    x_ticks = np.arange(0, 24.5, 1)

    # 第一子图：BGMEDIAN & 3σ 过滤
    ax1 = plt.subplot(5, 1, 1)
    ax1.plot(time_hours, bg_medianr, label="BGMEDIAN (Raw)", color="blue", linestyle="-")
    ax2 = ax1.twinx()
    ax2.plot(time_hours, bg_median_filteredr, label="BGMEDIAN<800", color="red", linestyle="-")
    ax1.set_ylabel("BGMEDIAN")
    ax2.set_ylabel("BGMEDIAN<800")
    ax1.set_title(f"BGMEDIAN of R-Band Over Time on {date_str}")
    ax1.grid(True)
    ax1.set_xticks(x_ticks)
    ax1.legend(loc="upper left")
    ax2.legend(loc="upper right")
    

    ax1 = plt.subplot(5, 1, 2)
    ax1.plot(time_hours, bg_medianb, label="BGMEDIAN (Raw)", color="blue", linestyle="-")
    ax2 = ax1.twinx()
    ax2.plot(time_hours, bg_median_filteredb, label="BGMEDIAN<800", color="red", linestyle="-")
    ax1.set_ylabel("BGMEDIAN")
    ax2.set_ylabel("BGMEDIAN<800")
    ax1.set_title(f"BGMEDIAN of B-Band Over Time on {date_str}")
    ax1.grid(True)
    ax1.set_xticks(x_ticks)
    ax1.legend(loc="upper left")
    ax2.legend(loc="upper right")

    # 第三子图：EE70 & 3σ 过滤
    ax3 = plt.subplot(5, 1, 3)
    ax3.plot(time_hours, ee70r, label="EE70 (Raw)", color="blue", linestyle="-")
    ax4 = ax3.twinx()
    ax4.plot(time_hours, ee70_filteredr, label="EE70<3", color="red", linestyle="-")
    ax3.set_ylabel("EE70")
    ax4.set_ylabel("EE70<3")
    ax3.set_title(f"EE70 of R-Band Over Time on {date_str}")
    ax3.grid(True)
    ax3.set_xticks(x_ticks)
    ax3.legend(loc="upper left")
    ax4.legend(loc="upper right")
    
    ax3 = plt.subplot(5, 1, 4)
    ax3.plot(time_hours, ee70b, label="EE70 (Raw)", color="blue", linestyle="-")
    ax4 = ax3.twinx()
    ax4.plot(time_hours, ee70_filteredb, label="EE70<4", color="red", linestyle="-")
    ax3.set_ylabel("EE70")
    ax4.set_ylabel("EE70<4")
    ax3.set_title(f"EE70 of B-Band Over Time on {date_str}")
    ax3.grid(True)
    ax3.set_xticks(x_ticks)
    ax3.legend(loc="upper left")
    ax4.legend(loc="upper right")
    
    ax3 = plt.subplot(5, 1, 5)
    ax3.plot(time_hours, sun_angs, label="sun_angs", color="red", linestyle="-")
    ax4 = ax3.twinx()
    ax4.plot(time_hours, moon_angs, label="moon_angs", color="blue", linestyle="-")
    ax3.set_ylabel("sun_angs")
    ax4.set_ylabel("moon_angs")
    ax3.set_title(f"sun and moon angle Over Time on {date_str}")
    ax3.grid(True)
    ax3.set_xticks(x_ticks)
    ax3.legend(loc="upper left")
    ax4.legend(loc="upper right")

    plt.tight_layout()
    plt.savefig(os.path.join(save_dir, f"{date_str}_plot.png"))
    plt.close()
    
def plotOneDay(json_file, srcPath, saveDir):
    
    date_str = json_file.replace(".json", "")
    json_path = os.path.join(srcPath, json_file)
    print(f"Processing {json_file}...")

    time_hoursr, bg_medianr, bg_stdr, ee70r, ee70_rmsr, sun_angsr, moon_angsr = parse_json(json_path, bandType="R")
    time_hoursb, bg_medianb, bg_stdb, ee70b, ee70_rmsb, sun_angsb, moon_angsb = parse_json(json_path, bandType="B")

    if len(time_hoursr) > 0:
        plot_data(saveDir, date_str, time_hoursr, bg_medianr, ee70r, bg_medianb, ee70b, sun_angsr, moon_angsr)
        print(f"Saved plot for {date_str}")
    else:
        print(f"No valid data for {date_str}, skipping.")

# 处理所有 JSON 文件
def process_all_json():
    
    # JSON 文件目录
    JSON_OUTPUT_DIR = "/data/vxpp_info/L1X/imgInfo"
    IMG_CURVE_OUTPUT_DIR = "/data/vxpp_info/L1X/imgInfo_curve"

    for year in sorted(os.listdir(JSON_OUTPUT_DIR), reverse=True):
        year_dir = os.path.join(JSON_OUTPUT_DIR, year)
        save_dir = os.path.join(IMG_CURVE_OUTPUT_DIR, year)
        if not os.path.isdir(year_dir):
            continue

        for json_file in sorted(os.listdir(year_dir), reverse=True):
            try:
                if not json_file.endswith(".json"):
                    continue
                
                # if json_file != "2025-01-01.json":
                #     continue

                date_str = json_file.replace(".json", "")
                saveImgName = f"{date_str}_plot.png"
                full_save_dir = os.path.join(save_dir, saveImgName)
                if os.path.exists(full_save_dir):
                    print(f"Skipping {json_file} because it already exists.")
                    continue
                
                plotOneDay(json_file, year_dir, save_dir)
                # break
            except Exception as e:
                print(f"Error processing {json_file}: {e}")
        # break

# 运行主程序
if __name__ == "__main__":
    process_all_json()
